from PyPDF2 import PdfFileWriter, PdfFileReader
import sys
import time

def extract_pdf(target_path, pagesNum):
    """
    抓取对应的pdf文件页面合成一个文件
    :param target_path: 抓取文件的路径
    :param pagesNum:    抓取的页面，格式[start-end, pageNum, pageNum, start-end……]
    :return:            None
    """
    reader = PdfFileReader(target_path)
    writer = PdfFileWriter()

    pages = []
    for des in pagesNum:
        if des.__contains__("-"):
            start = int(des.split("-")[0])
            end = int(des.split("-")[1])
            for i in range(start, end+1):
                pages.append(i)
        else:
            pages.append(int(des))

    getpages = list()

    for i in pages:
        page = reader.getPage(i - 1)
        getpages.append(page)

    for page in getpages:
        writer.addPage(page)

    output_filename = "./extract_" + str(int(time.time())) +".pdf"
    outputStream = open(output_filename, "wb")
    writer.write(outputStream)
    outputStream.close()


if __name__ == '__main__':
    if len(sys.argv) == 1:
        print("Error：请正确输入参数 ")
        print("提取文件所在路径，提取页数……[格式：start-end, pageNum, pageNum, start-end……]")
    else:
        extract_pdf(sys.argv[1], sys.argv[2:])
